"""
  此文件从过滤的数据中提取出换道数据
  filter_datas: 三维数组
  filter_data: 二维数组（三位数组内部的）
  data: 一维数组（二维数组内部的）

  因为内部的复杂校验逻辑，所以此文件的运行时常可能会比较久
"""

# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)

# 导入内部包
from utils.read_data import read_original_data
from utils.write_data import write_data_to_file

# 文件路径常量
FILE_PATH_I80_1_from = "../../Data/Ngsim数据集/I80数据集/2. 过滤数据/trajectories-0400-0415_filter.txt" # 2052 884 18
FILE_PATH_I80_2_from = "../../Data/Ngsim数据集/I80数据集/2. 过滤数据/trajectories-0500-0515_filter.txt" # 1836 526 18
FILE_PATH_I80_3_from = "../../Data/Ngsim数据集/I80数据集/2. 过滤数据/trajectories-0515-0530_filter.txt" # 1790 1153 18
FILE_PATH_101_1_from = "../../Data/Ngsim数据集/101数据集/2. 过滤数据/trajectories-0750am-0805am_filter.txt" # 2169 437 18
FILE_PATH_101_2_from = "../../Data/Ngsim数据集/101数据集/2. 过滤数据/trajectories-0805am-0820am_filter.txt" # 2017 569 18
FILE_PATH_101_3_from = "../../Data/Ngsim数据集/101数据集/2. 过滤数据/trajectories-0820am-0835am_filter.txt" # 1915 421 18

FILE_PATH_I80_1_to = "../../Data/Ngsim数据集/I80数据集/3. 提取数据/2. 换道数据/trajectories-0400-0415_lanechange.txt" # 300 80 21 / 过滤非法数据后 91 80 24（师兄原始的实验数据量 260 297 14）
FILE_PATH_I80_2_to = "../../Data/Ngsim数据集/I80数据集/3. 提取数据/2. 换道数据/trajectories-0500-0515_lanechange.txt" # 251 80 21 / 过滤非法数据后 71 80 24（师兄原始的实验数据量 212 297 14）
FILE_PATH_I80_3_to = "../../Data/Ngsim数据集/I80数据集/3. 提取数据/2. 换道数据/trajectories-0515-0530_lanechange.txt" # 281 80 21 / 过滤非法数据后 73 80 24（师兄原始的实验数据量 211 297 14）
FILE_PATH_101_1_to = "../../Data/Ngsim数据集/101数据集/3. 提取数据/2. 换道数据/trajectories-0750am-0805am_lanechange.txt" # 293 80 21 / 过滤非法数据后 148 80 24（师兄原始的实验数据量 215 297 14）
FILE_PATH_101_2_to = "../../Data/Ngsim数据集/101数据集/3. 提取数据/2. 换道数据/trajectories-0805am-0820am_lanechange.txt" # 300 80 21 / 过滤非法数据后 163 80 24（师兄原始的实验数据量 242 297 14）
FILE_PATH_101_3_to = "../../Data/Ngsim数据集/101数据集/3. 提取数据/2. 换道数据/trajectories-0820am-0835am_lanechange.txt" # 326 80 21 / 过滤非法数据后 136 80 24（师兄原始的实验数据量 227 297 14）

# 其他常量
TARGET_FRAMES = 80 # 目标帧数（要收集的数据数量）
BEFORE_LANECHANGE_STABLE_FRAMES = 40 # 换道前的稳定运行帧数
AFTER_LANECHANGE_STABLE_FRAMES = 40 # 换道后的稳定运行帧数
EFFECTIVE_FOLLOW_SPACE = 120 # 有效的跟驰距离（当超过这个距离时，不认为有跟驰效果）

loopNum = 0

# 提取换道数据
def extract_lanechange_data(filter_datas, lane_idx=13, front_carId_idx=14, after_carId_idx=15):
  # 收集到的若干换道数据组
  lanechange_datas = []

  for i in range(len(filter_datas)):
    filter_data = filter_datas[i]

    # 自车上一帧的车道号
    pre_laneId = filter_data[0][lane_idx]

    for j in range(len(filter_data)):
      # !因为数据量，所以这个循环需要执行百万次！！！
      global loopNum
      loopNum += 1
      print(loopNum)

      # 读取数据（当前 车道Id）
      cur_laneId = filter_data[j][lane_idx]

      if not is_lanechange(pre_laneId, cur_laneId) or not is_valid_lane(pre_laneId) or not is_valid_lane(cur_laneId):
        pre_laneId = cur_laneId
        continue
      # !到达这里，则自车发生了换道，且符合车道号条件

      if not has_around_car(filter_data[j], front_carId_idx, after_carId_idx):
        pre_laneId = cur_laneId
        continue
      # !到达这里，则自车有前后车辆

      front_carId = filter_data[j][front_carId_idx]
      after_carId = filter_data[j][after_carId_idx]
      # 找前、后车数据（基于车ID）
      front_car_data = find_data_by_carId(filter_datas, front_carId)
      after_car_data = find_data_by_carId(filter_datas, after_carId)

      has_enough_frames_before, front_car_data_start_idx, after_car_data_start_idx = has_enough_stable_frames_before_lanechange(filter_data, front_car_data, after_car_data, j - BEFORE_LANECHANGE_STABLE_FRAMES, pre_laneId, cur_laneId, lane_idx)
      has_enough_frames_after = has_enough_stable_frames_after_lanechange(filter_data, j, cur_laneId, front_carId, after_carId, lane_idx)
      if not has_enough_frames_before or not has_enough_frames_after:
        pre_laneId = cur_laneId
        continue
      # !到达这里，则自车有足够的换道前后稳定运行的数据帧

      # 收集换道数据
      lanechange_data = collect_lanechange_data(filter_data, front_car_data, after_car_data, j - BEFORE_LANECHANGE_STABLE_FRAMES, front_car_data_start_idx, after_car_data_start_idx)
      lanechange_datas.append(lanechange_data)

      # !尝试继续向后检测（有可能这辆自车后边又出现变道了）
      pre_laneId = cur_laneId

  return lanechange_datas

def collect_lanechange_data(self_car_data, front_car_data, after_car_data, self_car_data_start_idx, front_car_data_start_idx, after_car_data_start_idx, frameId_idx=1):
  # todo 目前的需求是：收集80帧的三辆车的运行数据。但是，目前会遇到一个问题：在这车从起始帧开始运行时，前车或后车(主要是后车)，可能不存在这个起始帧所对应的数据，应该在前边的has_enough_stable环节过滤掉这种数据。
  lanechange_data = []

  # !开始收集数据（三个数据组合并成一个数据组）
  for i in range(TARGET_FRAMES):
    if(front_car_data_start_idx + i >= len(front_car_data)):
      print("front数据越界", front_car_data_start_idx, i)
      print(front_car_data[0])

    if(after_car_data_start_idx + i >= len(after_car_data)):
      print("after数据越界", after_car_data_start_idx, i)
      print(after_car_data[0])

    front_car_item = front_car_data[front_car_data_start_idx + i]
    self_car_item = self_car_data[self_car_data_start_idx + i]
    after_car_item = after_car_data[after_car_data_start_idx + i]

    # 安全性校验
    if(front_car_item[frameId_idx] != self_car_item[frameId_idx] or self_car_item[frameId_idx] != after_car_item[frameId_idx]):
      print("数据帧不一致", front_car_item[frameId_idx], self_car_item[frameId_idx], after_car_item[frameId_idx])

    # !直接合并数据，并且做收集
    lanechange_data.append(format_lanechange_data(self_car_item, front_car_item, after_car_item))

  return lanechange_data

# 基于carId找数据
def find_data_by_carId(filter_datas, target_carId, carId_idx=0):
  for i in range(len(filter_datas)):
    if filter_datas[i][0][carId_idx] == target_carId:
      # 找到前车数据
      return filter_datas[i]
    
  print("无法基于车ID，找到对应的数据。出现bug了~")
  # return []

# 基于目标帧找起始索引
def find_start_idx_by_frame(filter_data, target_frame, frameId_idx=1):
  for j in range(len(filter_data)):
    if filter_data[j][frameId_idx] == target_frame:
      # 找到目标起始帧
      return j
    
  print("无法找到当前帧所对应的那条数据。出现bug了~")
  # return -1

# 对生成的换道数据做格式化
def format_lanechange_data(self_car_item, front_car_item, after_car_item):
  combine_data = []

  # 数据帧
  combine_data.append(self_car_item[1])

  # 编号
  combine_data.append(front_car_item[0])
  combine_data.append(self_car_item[0])
  combine_data.append(after_car_item[0])

  # X坐标（横轴）
  combine_data.append(front_car_item[4])
  combine_data.append(self_car_item[4])
  combine_data.append(after_car_item[4])

  # Y坐标（纵轴）
  combine_data.append(front_car_item[5])
  combine_data.append(self_car_item[5])
  combine_data.append(after_car_item[5])

  # 速度
  combine_data.append(front_car_item[11])
  combine_data.append(self_car_item[11])
  combine_data.append(after_car_item[11])

  # 加速度
  combine_data.append(front_car_item[12])
  combine_data.append(self_car_item[12])
  combine_data.append(after_car_item[12])

  # 车道编号（自车车道。因为自车从旧车道移动到新车道，所以自车的数据中包含了原始车道和目标车道）
  combine_data.append(self_car_item[13])

  # 车头间距（分别为自车与前车、后车与自车、前车与后车）
  combine_data.append(abs(self_car_item[5] - front_car_item[5])) # todo 这计算的对吗
  combine_data.append(abs(self_car_item[5] - after_car_item[5])) # todo 这计算的对吗
  combine_data.append(abs(front_car_item[5] - after_car_item[5])) # todo 这计算的对吗

  # 横向间距
  combine_data.append(abs(self_car_item[4] - front_car_item[4]))
  combine_data.append(abs(self_car_item[4] - after_car_item[4]))

  # 车头时距（分别为自车与前车、后车与自车）
  # todo 这里的车头时距应该重新基于位置和速度来算
  combine_data.append(self_car_item[17])
  combine_data.append(after_car_item[17])

  return combine_data


# 是否发生换换道
def is_lanechange(pre_laneId, cur_laneId):
  if pre_laneId == cur_laneId:
    return False
  
  return True

# 是否在合法车道内
def is_valid_lane(cur_laneId):
  return cur_laneId >= 1 and cur_laneId <= 5

# 自车是否有周围车辆
def has_around_car(data, front_carId_idx=14, after_carId_idx=15):
  front_carId = data[front_carId_idx]
  after_carId = data[after_carId_idx]

  if front_carId == 0 or after_carId == 0:
    return False
  
  return True

# 是否有足够的换道前的运行帧数（车道无变化）
def has_enough_stable_frames_before_lanechange(self_car_data, front_car_data, after_car_data, start_idx, pre_laneId, cur_laneId, lane_idx, frameId_idx=1, total_frames_idx=2):
  # 边界校验
  if start_idx < 0 or start_idx + BEFORE_LANECHANGE_STABLE_FRAMES > len(self_car_data):
    return False, -1, -1
    
  # 帧数量和范围校验（为了防止自车有起始帧和结束帧，但是前车或后车没有对应的起始帧或结束帧）
  ## 计算帧
  self_car_start_frame = self_car_data[start_idx][frameId_idx]
  front_car_start_frame = front_car_data[0][frameId_idx]
  after_car_start_frame = after_car_data[0][frameId_idx]
  front_car_total_frame = front_car_data[0][total_frames_idx]
  after_car_total_frame = after_car_data[0][total_frames_idx]
  ## 校验
  if front_car_start_frame > self_car_start_frame or after_car_start_frame > self_car_start_frame:
    return False, -1, -1
  if front_car_start_frame + front_car_total_frame < self_car_start_frame + TARGET_FRAMES or after_car_start_frame + after_car_total_frame < self_car_start_frame + TARGET_FRAMES:
    return False, -1, -1
  
  # !找前车数据的起始索引（基于自车起始帧）
  front_car_data_start_idx = find_start_idx_by_frame(front_car_data, self_car_start_frame, frameId_idx)
  if not front_car_data_start_idx and front_car_data_start_idx != 0:
    print("无法找到当前帧所对应的前车数据。出现bug了~")

  # !找后车数据的起始索引（基于后车起始帧）
  after_car_data_start_idx = find_start_idx_by_frame(after_car_data, self_car_start_frame, frameId_idx)
  if not after_car_data_start_idx and after_car_data_start_idx != 0:
    print("无法找到当前帧所对应的后车数据。出现bug了~")
    
  # 车道号校验（自车 和 前车、后车）
  for i in range(BEFORE_LANECHANGE_STABLE_FRAMES):
    if self_car_data[start_idx + i][lane_idx] != pre_laneId:
      return False, -1, -1
    if front_car_data[front_car_data_start_idx + i][lane_idx] != cur_laneId:
      return False, -1, -1
    if after_car_data[after_car_data_start_idx + i][lane_idx] != cur_laneId:
      return False, -1, -1
    
  return True, front_car_data_start_idx, after_car_data_start_idx
  

# 是否有足够的换道后的稳定运行帧数（稳定，指的是前后车ID不变）
def has_enough_stable_frames_after_lanechange(filter_data, start_idx, cur_laneId, target_front_curId, target_after_carId, lane_idx, front_carId_idx=14, after_carId_idx=15):
  # 边界校验
  if start_idx < 0 or start_idx + AFTER_LANECHANGE_STABLE_FRAMES > len(filter_data):
    return False
  
  # 车道号、周车数、前车ID、后车ID 校验
  for i in range(AFTER_LANECHANGE_STABLE_FRAMES):
    item_idx = start_idx + i

    cur_front_curId = filter_data[item_idx][front_carId_idx]
    cur_after_curId = filter_data[item_idx][after_carId_idx]

    if filter_data[item_idx][lane_idx] != cur_laneId or not has_around_car(filter_data[item_idx]) or (cur_front_curId != target_front_curId) or (cur_after_curId != target_after_carId):
      return False
    
  return True

# 过滤掉有问题的组（例如，前后车位置重叠等）
def filter_error_group(lanechange_datas, front_car_pos_idx=7, self_car_pos_idx=8, after_car_pos_idx=9, front_car_v_idx=10, self_car_v_idx=11, after_car_v_idx=12):
  filtered_datas = []

  has_error_data = False

  for i, lanechange_data in enumerate(lanechange_datas):
    # !检测这组数据是否有问题（从换道后开始检测，在车道变换前不做要求）
    for j in range(AFTER_LANECHANGE_STABLE_FRAMES, TARGET_FRAMES):
      data = lanechange_data[j]

      if not is_valid_follow_data(data, front_car_pos_idx, self_car_pos_idx, after_car_pos_idx, front_car_v_idx, self_car_v_idx, after_car_v_idx):
        print("出现一条不符合规则的跟车数据", i, j)
        has_error_data = True
        break

    if not has_error_data:
      # 收集这组数据
      filtered_datas.append(lanechange_data)

    has_error_data = False

  return filtered_datas

def is_valid_follow_data(data, front_car_pos_idx, self_car_pos_idx, after_car_pos_idx, front_car_v_idx, self_car_v_idx, after_car_v_idx):
  front_and_self_space = data[front_car_pos_idx] - data[self_car_pos_idx]
  self_and_after_space = data[self_car_pos_idx] - data[after_car_pos_idx]
  front_car_v = data[front_car_v_idx]
  self_car_v = data[self_car_v_idx]
  after_car_v = data[after_car_v_idx]

  # 限制车头间距（这里的0可以再优化，例如设置成可接受的最小间距）# todo 这个部分的 EFFECTIVE_FOLLOW_SPACE 设置成了定值，但是我感觉可以基于车速来动态生成（可以后续做优化）
  if front_and_self_space <= 0 or self_and_after_space <= 0:
    print("两车间距出现负值", front_and_self_space, self_and_after_space)
    return False
  if front_and_self_space > EFFECTIVE_FOLLOW_SPACE or self_and_after_space > EFFECTIVE_FOLLOW_SPACE:
    print("车头间距超出阈值", front_and_self_space, self_and_after_space)
    return False
  if front_car_v <= 0 or self_car_v <= 0 or after_car_v <= 0:
    print("车速出现负值", front_car_v, self_car_v, after_car_v)
    return False

  return True

def handle(fromFileName, toFileName, index):
  # 读取数据
  filter_datas = read_original_data(fromFileName)

  # 提取换道数据
  lanechange_datas = extract_lanechange_data(filter_datas)

  # 过滤出含非法值的数据组
  filtered_datas = filter_error_group(lanechange_datas)

  print("第几组数据", index + 1)
  print("提取结果", len(lanechange_datas), len(lanechange_datas[0]), len(lanechange_datas[0][0]))
  print("过滤后的结果", len(filtered_datas), len(filtered_datas[0]), len(filtered_datas[0][0]))

  # 写入数据
  write_data_to_file(filtered_datas, toFileName)

if __name__ == "__main__":
  fileNames = [
    [FILE_PATH_I80_1_from, FILE_PATH_I80_1_to],
    [FILE_PATH_I80_2_from, FILE_PATH_I80_2_to],
    [FILE_PATH_I80_3_from, FILE_PATH_I80_3_to],
    [FILE_PATH_101_1_from, FILE_PATH_101_1_to],
    [FILE_PATH_101_2_from, FILE_PATH_101_2_to],
    [FILE_PATH_101_3_from, FILE_PATH_101_3_to]
  ]

  for i in range(len(fileNames)):
    fileName = fileNames[i]
    handle(fileName[0], fileName[1], i)